import os
import hou
import time

try:
    import pathlib
    import getpass
    import CONFIG as cfg
    from Disk import getDiskFreeSpace

    deadline = cfg.MZ_CONFIG.get('deadline', {})
except Exception as err:
    pass

i_list = ['FX_mantra_','arnold']
render_node_list = ['fx_abc', 'FX_rop_geometry_']
work_nodes = ['FX_mantra_', 'fx_abc', 'FX_rop_geometry_','arnold']



# 节点检测
def checkNodes(check_type):
    right_nodes = True
    wrong_nodes = []
    if check_type == 'Render':
        num = hou.parm('num').evalAsInt()
        if num:
            for n in range(num):
                node_path = hou.parm('node{}'.format(n + 1)).evalAsString()
                if node_path:
                    node = hou.node(node_path)
                    if node.type().name() in work_nodes:
                        pass
                    else:
                        wrong_nodes.append(node_path)
    if check_type == 'Submit':
        job = hou.parm('job').evalAsInt()
        if job:
            for j in range(job):
                n_path = hou.parm('submitnode{}'.format(j + 1)).evalAsString()
                if n_path:
                    s_node = hou.node(n_path)
                    if s_node.type().name() in work_nodes:
                        pass
                    else:
                        wrong_nodes.append(n_path)
    if wrong_nodes:
        wrong_message = '\n'.join(wrong_nodes)
        right_nodes = False
        hou.ui.displayMessage('错误的节点类型:\n' + wrong_message)
    return right_nodes


# GO 顺序解算
def createRenderScript(args):
    direct_render, interval_time, current_node_path, next_node_path = args
    script = """import MZ_tool as mz\nmz.interval_render({},{},'{}','{}')""".format(direct_render, interval_time,
                                                                                    current_node_path, next_node_path)
    return script


def Go(*args):
    check = checkNodes('Render')
    if not check:
        return
    cn = hou.pwd()
    num = cn.parm('num').evalAsInt()

    # 获取节点
    node_list = []
    if num:
        for i in range(num):
            node_path = hou.parm('node{}'.format(i + 1)).eval()
            if node_path:
                node = hou.node(node_path)
                name = node.type().name()
                if name in i_list or name in render_node_list:
                    node_list.append(node)
                    node.parm('postrender').set('')
                    node.parm('lpostrender').set('python')

    if node_list:
        node_script_args_list = []
        interval_time = hou.parm('itime').evalAsInt()
        for i in range(len(node_list)):
            node = node_list[i]
            node_path = node.path()
            direct_render = 1
            if node.type().name() in render_node_list:
                direct_render = 0
            node_args = [direct_render, interval_time, node_path]
            node_script_args_list.append(node_args)
        for s in range(len(node_list)):
            node = node_list[s]

            if s + 1 < len(node_list):
                a_s = node_script_args_list[s]
                next_node_path = node_list[s + 1].path()
                a_s.append(next_node_path)
                script = createRenderScript(a_s)
                node.parm('postrender').set(script)
                node.parm('lpostrender').set('python')

        first_node = node_list[0]
        first_node_direct = node_script_args_list[0][0]
        # print('first_node_direct:',first_node_direct)
        if first_node_direct:
            # print(first_node)
            first_node.render()
        else:
            # print(first_node)
            first_node.node('render').render()


# 设置解算优先级

def setSim(*args):
    num = hou.parm('job').evalAsInt()
    if num:
        for i in range(num):
            sim = hou.parm('sim{}'.format(i + 1)).evalAsInt()
            if sim:
                hou.parm('chunk{}'.format(i + 1)).set(99999)
                hou.parm('priority{}'.format(i + 1)).set(90)
            else:
                priority = hou.parm('priority{}'.format(i + 1)).evalAsInt()
                if priority > 50:
                    hou.parm('chunk{}'.format(i + 1)).set(10)
                    hou.parm('priority{}'.format(i + 1)).set(50)


class DeadlineTool(object):

    def __init__(self):
        super(DeadlineTool, self).__init__()
        self.origin_node_outputpath = {}
        self.concurrent_filter = ['FX_mantra_','arnold']  # 用来添加多开渲染机器

    def add_deadline_outpath(self, export_type, node_path, raw_outputpath, date_time, var):
        """
        将输出路径转成deadline的输出路径
        """
        if raw_outputpath:
            job_temp = '$JOB/img' if export_type == 'img' else '$JOB/geo'
            if node_path not in self.origin_node_outputpath.keys():  # 当不存在在原始节点时，存入原始节点输出路径
                self.origin_node_outputpath[node_path] = raw_outputpath
                # dir_path,filename=os.path.split(outputpath)
                # outputpath= raw_outputpath.replace(job_temp,'`@job`/{}/{}/wedge{}'.format(export_type,date_time,var))
                outputpath = raw_outputpath.replace(job_temp, '`@job`/{}'.format(export_type))

            else:
                # outputpath= self.origin_node_outputpath[node_path].replace(job_temp,'`@job`/{}/{}/wedge{}'.format(export_type,date_time,var))
                outputpath = self.origin_node_outputpath[node_path].replace(job_temp, '`@job`/{}'.format(export_type))
                # dir_path,filename=os.path.split(self.origin_node_outputpath[node_path])
            # outputpath=os.path.join(dir_path,'deadline',date_time,'var{}'.format(var),filename)

        return outputpath  # .replace('\\','/')

    def set_deadline_outputpath(self, node, date_time, var):
        """
        将输出路径转成deadline的输出路径
        """
        if node.type().name() == 'rop_comp':
            outputpath = node.parm('coppath').rawValue()
            node.parm('coppath').set(self.add_deadline_outpath('img', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'rop_alembic':
            outputpath = node.parm('filename').rawValue()
            node.parm('filename').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'rop_geometry':
            outputpath = node.parm('sopoutput').rawValue()
            node.parm('sopoutput').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'rop_fbx':
            outputpath = node.parm('sopoutput').rawValue()
            node.parm('sopoutput').set(self.add_deadline_outpath('fbx', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'rop_gltf':
            outputpath = node.parm('file').rawValue()
            node.parm('file').set(self.add_deadline_outpath('gltf', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'FX_mantra_':
            outputpath = node.parm('vm_picture').rawValue()
            node.parm('vm_picture').set(self.add_deadline_outpath('img', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'ifd':
            outputpath = node.parm('vm_picture').rawValue()
            node.parm('vm_picture').set(self.add_deadline_outpath('img', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'filecache':
            outputpath = node.parm('file').rawValue()
            node.parm('file').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'FX_rop_geometry_':
            outputpath = node.parm('sopoutput').rawValue()
            node.parm('sopoutput').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'mz_recache':
            outputpath = node.parm('file').rawValue()
            node.parm('file').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))
        elif node.type().name() == 'mz_DopImport':
            outputpath = node.parm('file').rawValue()
            node.parm('file').set(self.add_deadline_outpath('geo', node.path(), outputpath, date_time, var))

    def init_var(self):
        var = hou.contextOption('wedge')
        if var and int(float(var)) >= 0:
            pass
        else:
            hou.setContextOption('wedge', -1)

    def checkNode(self, node):
        """
        检查传入的节点类型
        """

        if node:
            if node.type().name() in i_list:
                return True, 0
            elif node.type().name() in render_node_list:
                return True, 1
            else:
                return False, 0
        else:
            return False, 0

    def submitDeadlineOrder(self, submit_time, var, deadline_hip_path):
        """
        排序提交
        """
        try:
            import DeadlineConnect as Connect
        except Exception as DeadlineConnect_ERR:
            print('DeadlineConnect_ERR:', str(DeadlineConnect_ERR))
            return
        if not deadline:
            return
        connectionObject = Connect.DeadlineCon(deadline['web_serive'], deadline['port'])
        file_path = deadline_hip_path
        version = hou.applicationVersionString()
        hip_name = os.path.splitext(os.path.split(file_path)[1])[0]
        user_name = getpass.getuser()

        cn = hou.pwd()
        num = cn.parm('job').evalAsInt()

        # 获取该次提交的所有待提交任务内容，通过order作为key
        all_jobs = {}  # {1:[job_data1,job_data2],2:[...],...}

        if num:
            for i in range(num):
                node_path = hou.parm('submitnode{}'.format(i + 1)).eval()
                node = hou.node(node_path) if node_path else None
                if not node:  # 当为空节点时
                    pass
                else:
                    # canRender,rtype=self.checkNode(node)
                    # if not canRender:
                    #     hou.ui.displayMessage('{} can not render !'.format(node_path))
                    #     return
                    node_name = node.name()
                    f1 = node.parm('f1').evalAsString()
                    f2 = node.parm('f2').evalAsString()
                    order = hou.parm('order{}'.format(i + 1)).evalAsInt()
                    pool = hou.parm('pool{}'.format(i + 1)).evalAsString()
                    chunk = hou.parm('chunk{}'.format(i + 1)).evalAsInt()
                    sim = hou.parm('sim{}'.format(i + 1)).evalAsInt()
                    pri = hou.parm('priority{}'.format(i + 1)).evalAsInt()
                    concurrentTasks = hou.parm('concurrentTasks{}'.format(i + 1)).evalAsInt()

                    if pri > 50 and not sim:
                        pri = 50

                    # 转deadline输出路径
                    self.set_deadline_outputpath(node, submit_time, var)

                    # 计算job
                    props = {}
                    props['Name'] = hip_name + '_' + node_name
                    props['UserName'] = user_name
                    props['Frames'] = '{}-{}'.format(f1, f2)
                    props['Plugin'] = 'Houdini'
                    props['Pool'] = pool
                    props['Priority'] = pri
                    props['ChunkSize'] = chunk
                    # if node.type().name() in self.concurrent_filter:
                    if not sim: # by is simlation
                        props['ConcurrentTasks'] = concurrentTasks

                    # 插件信息
                    plug = {"OutputDriver": node_path,
                            "SceneFile": deadline_hip_path,
                            "Version": version
                            }

                    # OutputDriver验证修改
                    if node.type().name() in render_node_list:
                        plug["OutputDriver"] = plug["OutputDriver"] + '/render'

                    job_data = {}
                    job_data['props'] = props
                    job_data['plug'] = plug

                    if order not in all_jobs.keys():
                        all_jobs[order] = [job_data]
                    else:
                        all_jobs[order].append(job_data)

        # 提交job
        if all_jobs:
            order_list = sorted(all_jobs.keys())
            pre_id_list = []
            for o in order_list:
                job_list = all_jobs[o]
                order_id_list = []  # 该次order提交后的job id列表
                for data in job_list:
                    if pre_id_list:  # 当存在每order的提交id时,增加依赖
                        data['props']['JobDependencies'] = ','.join(pre_id_list)
                        # 提交job
                        job = connectionObject.Jobs.SubmitJob(data['props'], data['plug'])
                        job_id = job['_id']
                        order_id_list.append(job_id)
                    else:  # 当不存在预先的id列表时,即排序第一的任务，直接提交
                        job = connectionObject.Jobs.SubmitJob(data['props'], data['plug'])
                        job_id = job['_id']
                        order_id_list.append(job_id)

                pre_id_list = order_id_list  # 清掉预先的id，替换为当前order提交后的所有job id列表

    def before_scene_saved(self, new_fp):
        """
        不需要了，job已经预置了。
        """
        dir_path = os.path.dirname(new_fp)
        if dir_path.endswith('deadline'):
            job_path = os.path.dirname(os.path.dirname(dir_path)).replace('\\', '/')
            hou.setContextOption('job', job_path)
        else:
            job_path = os.path.dirname(dir_path).replace('\\', '/')
            hou.setContextOption('job', job_path)
            # current_path=hou.getenv('HIP')
        # print('current_path:',current_path)
        # if current_path.endswith('deadline'):
        #     job_path=os.path.dirname(os.path.dirname(current_path))
        #     #job_path=os.path.split(current_path)[0]
        #     hou
        # elif current_path.endswith('hip'):
        #     job_path=os.path.dirname(current_path)
        # else:

        #     job_path=os.path.dirname(current_path)
        # print('job_path:',job_path)
        # hou.putenv('JOB',job_path)

    def submit_Deadline(self, submit_time, var, deadline_hip_path):
        """
        只有前一个提交了之后才会有job id
        """

        try:
            import DeadlineConnect as Connect
        except Exception as DeadlineConnect_ERR:
            print('DeadlineConnect_ERR:', str(DeadlineConnect_ERR))
            return
        if not deadline:
            return
        connectionObject = Connect.DeadlineCon(deadline['web_serive'], deadline['port'])
        file_path = deadline_hip_path  # hou.hipFile.name()
        version = hou.applicationVersionString()
        hip_name = os.path.splitext(os.path.split(file_path)[1])[0]
        user_name = getpass.getuser()

        cn = hou.pwd()
        num = cn.parm('gjob').evalAsInt()

        pre_job_id = ''
        if num:
            for i in range(num):
                gnode_path = hou.parm('gnode{}'.format(i + 1)).eval()
                gnode = hou.node(gnode_path) if gnode_path else None
                if not gnode:  # 当为空节点时
                    pass
                else:
                    # canRender,rtype=self.checkNode(gnode)
                    # if not canRender:
                    #     hou.ui.displayMessage('{} can not render !'.format(gnode_path))
                    #     return
                    gnode_name = gnode.name()
                    f1 = gnode.parm('f1').evalAsString()
                    f2 = gnode.parm('f2').evalAsString()
                    gpool = hou.parm('gpool{}'.format(i + 1)).evalAsString()
                    gchunk = hou.parm('gchunk{}'.format(i + 1)).evalAsInt()
                    # gsim = hou.parm('gsim{}'.format(i + 1)).evalAsInt()  # group暂时不用，没有添加gim
                    gpri = hou.parm('gpriority{}'.format(i + 1)).evalAsInt()
                    childnum = hou.parm('child{}'.format(i + 1)).evalAsInt()
                    # gconcurrentTasks = hou.parm('gconcurrentTasks{}'.format(i + 1)).evalAsInt()
                    # 转deadline输出路径
                    self.set_deadline_outputpath(gnode, submit_time, var)

                    # 计算job
                    gprops = {}
                    gprops['Name'] = hip_name + '_' + gnode_name
                    gprops['UserName'] = user_name
                    gprops['Frames'] = '{}-{}'.format(f1, f2)
                    gprops['Plugin'] = 'Houdini'
                    gprops['Pool'] = gpool
                    gprops['Priority'] = gpri
                    gprops['ChunkSize'] = gchunk


                    # 插件信息
                    gplug = {"OutputDriver": gnode_path,
                             "SceneFile": deadline_hip_path,
                             "Version": version
                             }

                    # OutputDriver验证修改
                    if gnode.type().name() in render_node_list:
                        gplug["OutputDriver"] = gplug["OutputDriver"] + '/render'

                    # 提交job
                    gjob = connectionObject.Jobs.SubmitJob(gprops, gplug)
                    print('提交:', gjob)
                    pre_job_id = gjob['_id']

                    # 当存在关联子节点时,排序子节点任务
                    if childnum:
                        for c in range(childnum):
                            gcnode_path = hou.parm('gcnode{}_{}'.format(i + 1, c + 1)).eval()
                            gcnode = hou.node(gcnode_path) if gcnode_path else None
                            if not gcnode:  # 当为空子节点时
                                pass
                            else:
                                # canRender,rtype=self.checkNode(gcnode)
                                # if not canRender:
                                #     hou.ui.displayMessage('next node : {} can not render !'.format(gcnode_path))
                                #     return
                                gcnode_name = gcnode.name()
                                gcf1 = gcnode.parm('f1').evalAsString()
                                gcf2 = gcnode.parm('f2').evalAsString()
                                gcpool = hou.parm('gcpool{}_{}'.format(i + 1, c + 1)).evalAsString()
                                gcchunk = hou.parm('gcchunk{}_{}'.format(i + 1, c + 1)).evalAsInt()
                                gcpri = hou.parm('gcpriority{}_{}'.format(i + 1, c + 1)).evalAsInt()
                                gcconcurrentTasks = hou.parm('gcconcurrentTasks{}'.format(i + 1)).evalAsInt()
                                # 转deadline输出路径
                                self.set_deadline_outputpath(gcnode, submit_time, var)

                                # 计算子节点job
                                gcprops = {}
                                gcprops['Name'] = hip_name + '_' + gcnode_name
                                gcprops['UserName'] = user_name
                                gcprops['Frames'] = '{}-{}'.format(gcf1, gcf2)
                                gcprops['Plugin'] = 'Houdini'
                                gcprops['Pool'] = gcpool
                                gcprops['Priority'] = gcpri
                                gcprops['ChunkSize'] = gcchunk
                                # if gcnode.type().name() in self.concurrent_filter:
                                #     gcprops['ConcurrentTasks'] = gcconcurrentTasks

                                # 插件信息
                                gcplug = {"OutputDriver": gcnode_path,
                                          "SceneFile": deadline_hip_path,
                                          "Version": version
                                          }

                                # OutputDriver验证修改
                                if gcnode.type().name() in render_node_list:
                                    gcplug["OutputDriver"] = gcplug["OutputDriver"] + '/render'

                                    # 生成子节点job
                                if not pre_job_id:
                                    gcjob = connectionObject.Jobs.SubmitJob(gcprops, gcplug)
                                    pre_job_id = gcjob['_id']
                                else:  # 添加依赖
                                    gcprops['JobDependencies'] = pre_job_id
                                    gcjob = connectionObject.Jobs.SubmitJob(gcprops, gcplug)
                                    pre_job_id = gcjob['_id']

                                print('\n')
                                print('子job属性:', gcprops)
                                print('\n')
                                print('gcjob:', gcjob)

                    pre_job_id = ''  # 一个组完成后，清空pre_job_id

        # hou.ui.displayMessage('提交成功')

    def Deadline(self, subtpye):
        result = hou.ui.displayMessage('请确认已保存工程!!', title='工程确认', buttons=('Submit', 'Save and Submit', 'Cancel'),
                                       default_choice=0)
        job_path = hou.contextOption('job')
        if job_path:
            hou.putenv('JOB', job_path)

        def submit():
            self.init_var()
            hip_path_list = []
            fn = hou.hipFile.name()
            bn = hou.hipFile.basename()

            pp = pathlib.Path(fn)
            dirname = pp.parent
            deadline = pathlib.Path(dirname, 'deadline')
            if not os.path.exists(deadline):
                os.makedirs(deadline)

            tar = pathlib.Path(deadline, bn)

            var_min = var_max = 0
            if hou.parm('wa').evalAsInt():
                var_min = hou.parm('varx').evalAsInt()
                var_max = hou.parm('vary').evalAsInt()

            submit_time = time.strftime('%Y-%m-%d__%H-%M-%S')
            for i in range(var_min, var_max + 1):
                stem = tar.stem
                # print(stem)  #cnn
                suf = tar.suffix
                # print(suf)   #.py
                deadline_hip = '{}_wedge{}_{}{}'.format(stem, i, submit_time, suf)
                new_fp = str(pathlib.Path(deadline, deadline_hip)).replace('\\', '/')  # 获取提供deadline的hip文件路径
                print(new_fp)
                hou.setContextOption('wedge', i)
                hip_path_list.append(new_fp)

                if subtpye == 'order':
                    self.submitDeadlineOrder(submit_time, i, new_fp)
                elif subtpye == 'group':
                    self.submit_Deadline(submit_time, i, new_fp)
                # self.before_scene_saved(new_fp)
                hou.hipFile.save(new_fp)

            freespace = self._getFreeSpace()
            result = '提交成功,\n 剩余空间{}T'.format(round(freespace, 2)) if freespace else '提交成功\n请映射{}盘符'.format(cfg.DISK)
            hou.ui.displayMessage(result)
            hou.hipFile.load(fn, suppress_save_prompt=True, ignore_load_warnings=True)

        if result == 0:
            submit()
        elif result == 1:
            hou.hipFile.save()
            submit()
            # print('hip_path_list:',hip_path_list)

    def _getFreeSpace(self):
        freespace = getDiskFreeSpace(cfg.DISK, 'T')
        return freespace


def submit_deadline_group(*args):
    dt = DeadlineTool()
    dt.Deadline('group')


def submit_deadline_order(*args):
    check = checkNodes('Submit')
    if not check:
        return
    dt = DeadlineTool()
    dt.Deadline('order')    